# Copyright (c) MONAI Consortium
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


from copy import deepcopy

# <placeholder> will be replaced by expert model outputs


expert_prefix = (
    "When answering the question, please analyze the image and "
    "incorporate the additional results generated by an expert classification model:\n"
)
text_only_expert_prefix = (
    "When answering the question, please incorporate the results generated by an expert classification model:\n"
)
multi_choice_prompt = [
    [
        0,
        "<image>\nThe following is a multiple-choice question about findings in chest X-ray in the frontal view. "
        f"{expert_prefix}<placeholder>"
        "Question: What are the potential abnormalities according to the provided X-ray image?\n"
        "<choices>"
        "Please reply with the corresponding answer choice letter(s).\n",
    ],
    [1, None],
]
binary_prompt = [
    [
        0,
        "<image>\nThe following is a question about findings in chest X-ray in the frontal view. "
        f"{expert_prefix}<placeholder>"
        "Question: is there <class_name> according to the provided X-ray image?\n"
        "Please reply with yes or no.\n",
    ],
    [1, None],
]

text_only_prompt = [
    [
        0,
        "The following is a question about findings in a chest X-ray in the frontal view. "
        f"{text_only_expert_prefix}<placeholder>"
        "Question: is there <class_name> according to the results?\n"
        "Please reply with yes or no.\n",
    ],
    [1, None],
]

model_list = (
    "Here is a list of available expert models:\n"
    "<BRATS(args)> Modality: MRI, Task: segmentation, Overview: A pre-trained model for volumetric (3D) "
    "segmentation of brain tumor subregions from multimodal MRIs based on BraTS 2018 data, "
    "Accuracy: Tumor core (TC): 0.8559 - Whole tumor (WT): 0.9026 - Enhancing tumor (ET): 0.7905 - Average: 0.8518, "
    "Valid args are: None\n"
    "<VISTA3D(args)> Modality: CT, Task: segmentation, Overview: domain-specialized interactive "
    "foundation model developed for segmenting and annotating human anatomies with precision, "
    "Accuracy: 127 organs: 0.792 Dice on average, Valid args are: 'everything', 'hepatic tumor',"
    "'pancreatic tumor', 'lung tumor', 'bone lesion', 'organs', 'cardiovascular', 'gastrointestinal', 'skeleton', "
    "or 'muscles'\n"
    "<VISTA2D(args)> Modality: cell imaging, Task: segmentation, Overview: model for cell segmentation, "
    "which was trained on a variety of cell imaging outputs, including brightfield, phase-contrast, "
    "fluorescence, confocal, or electron microscopy, Accuracy: Good accuracy across several cell imaging datasets, "
    "Valid args are: None\n"
    "<CXR(args)> Modality: chest x-ray (CXR), Task: classification, Overview: pre-trained model which are "
    "trained on large cohorts of data, Accuracy: Good accuracy across several diverse chest x-rays datasets, "
    "Valid args are: None\n"
    "Give the model <NAME(args)> when selecting a suitable expert model.\n"
)


def replace(prompts, placeholder, content):
    """Replace placeholder with content in prompts"""
    p = deepcopy(prompts)
    for idx, _ in enumerate(prompts):
        if prompts[idx][1] is None:
            continue
        p[idx][1] = prompts[idx][1].replace(placeholder, content)
    return p


def has_placeholder(prompts, placeholder):
    """Check if placeholder exists in prompts"""
    for _, words in prompts:
        if words is None:
            continue
        if placeholder in words:
            return True
    return False


_first_prompt = "Question: is there <class_name> according to the image?\n Please reply with yes or no.\n"
binary_conv_prompt = [
    [0, model_list + f"<image> This is a CXR image.\n{_first_prompt}"],
    [1, "This looks like a chest x-ray. Let me first trigger <CXR()>."],
    [
        0,
        f"The resulting predictions are:\n<placeholder>Analyze the image and take these predictions into "
        f"account when responding to this prompt.\n{_first_prompt}",
    ],
    [1, None],
]

templates = {
    "multi_choice": multi_choice_prompt,
    "binary": binary_prompt,
    "binary_conv": binary_conv_prompt,
    "text_only": text_only_prompt,
}
